package com.csust.automaticteach.manager.websocket;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.BinaryMessage;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.BinaryWebSocketHandler;

import java.nio.ByteBuffer;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Component
public class YjsWebSocketHandler extends BinaryWebSocketHandler {

    /** room -> sessions */
    private final Map<String, Set<WebSocketSession>> rooms = new ConcurrentHashMap<>();

    /* -------------------------------------------------- */

    @Override
    protected void handleBinaryMessage(WebSocketSession session,
                                       BinaryMessage message) throws Exception {
        String room = getRoom(session);
        broadcast(room, message.getPayload(), session);
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        String room = getRoom(session);
        rooms.computeIfAbsent(room, k -> ConcurrentHashMap.newKeySet())
                .add(session);
        log.info("加入房间 {}, 当前人数 {}", room, rooms.get(room).size());
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        String room = getRoom(session);
        Set<WebSocketSession> peers = rooms.get(room);
        if (peers != null) {
            peers.remove(session);
            log.info("离开房间 {}, 剩余人数 {}", room, peers.size());
        }
    }

    /* -------------------------------------------------- */

    /** 广播给同房间其他客户端（串行化发送，避免并发写异常） */
    private void broadcast(String room, ByteBuffer payload, WebSocketSession sender) {
        Set<WebSocketSession> peers = rooms.get(room);
        if (peers == null) return;

        peers.stream()
                .filter(WebSocketSession::isOpen)
                .filter(p -> !p.equals(sender))
                .forEach(p -> {
                    // 串行化写，防止 BINARY_PARTIAL_WRITING
                    synchronized (p) {
                        try {
                            p.sendMessage(new BinaryMessage(payload.duplicate()));
                        } catch (Exception e) {
                            log.warn("发送失败：{}", e.getMessage());
                        }
                    }
                });
    }

    /** 从 URI 提取 room */
    private String getRoom(WebSocketSession session) {
        String path = session.getUri().getPath();   // /yjs/123
        return path.substring(path.lastIndexOf('/') + 1);
    }
}